# ===----------------------------------------------------------------------=== #
# Copyright (c) 2025, Modular Inc. All rights reserved.
#
# Licensed under the Apache License v2.0 with LLVM Exceptions:
# https://llvm.org/LICENSE.txt
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===----------------------------------------------------------------------=== #

from memory import LegacyUnsafePointer as UnsafePointer
from math import ceildiv, isclose
from random import rand
from sys.info import simd_width_of

from itertools import product
from layout import IntTuple, LayoutTensor, RuntimeLayout, Layout
from nn.conv import ConvDirectNHWC, ConvInfoStatic, pack_filter
from nn.conv_utils import (
    ConvShape,
    get_direct_conv_micro_kernel_width,
    get_micro_kernel_shape,
)

from utils.index import Index, IndexList


fn test[
    N: Int,
    H: Int,
    W: Int,
    C: Int,
    R: Int,
    S: Int,
    F: Int,
    stride: IndexList[2],
    dilation: IndexList[2],
    pad_h: IndexList[2],
    pad_w: IndexList[2],
]() raises:
    # Output Shape.
    # fmt: off
    comptime HO = (H + pad_h[0] + pad_h[1] - dilation[1] * (R - 1) - 1) // stride[0] + 1
    comptime WO = (W + pad_w[0] + pad_w[1] - dilation[0] * (S - 1) - 1) // stride[1] + 1
    # fmt: on
    comptime type = DType.float32
    comptime simd_size = simd_width_of[type]()
    comptime num_groups = 1

    var conv_shape = ConvShape[2](
        n=N,
        input_dims=Index(H, W),
        output_dims=Index(HO, WO),
        filter_dims=Index(R, S),
        c=C,
        f=F,
        stride=stride,
        dilation=dilation,
        pad_d=Index(0, 0),
        pad_h=pad_h,
        pad_w=pad_w,
        num_groups=num_groups,
    )

    var input_ptr = UnsafePointer[Scalar[type]].alloc(N * H * W * C)
    var filter_ptr = UnsafePointer[Scalar[type]].alloc(R * S * C * F)

    # output from conv w/ dynamic and static shapes.
    var output_ptr_static = UnsafePointer[Scalar[type]].alloc(N * HO * WO * F)
    var output_ptr_dynamic = UnsafePointer[Scalar[type]].alloc(N * HO * WO * F)

    rand[type](input_ptr, N * H * W * C)
    rand[type](filter_ptr, R * S * C * F)

    comptime layout_4d = Layout.row_major[4]()
    comptime layout_5d = Layout.row_major[5]()
    var input = LayoutTensor[type, Layout.row_major(N, H, W, C)](input_ptr)
    var filter = LayoutTensor[type, layout_4d](
        filter_ptr, RuntimeLayout[layout_4d].row_major(Index(R, S, C, F))
    )
    var output_static = LayoutTensor[type, Layout.row_major(N, HO, WO, F)](
        output_ptr_static
    )
    var output_dynamic = LayoutTensor[type, layout_4d](
        output_ptr_dynamic,
        RuntimeLayout[layout_4d].row_major(Index(N, HO, WO, F)),
    )

    # Pre-packed filter for dynamic shapes.
    comptime micro_kernel_width_default = get_direct_conv_micro_kernel_width()
    comptime micro_kernel_f_size_default = micro_kernel_width_default * simd_size
    var rounded_F_dynamic = (
        ceildiv(F, micro_kernel_f_size_default) * micro_kernel_f_size_default
    )
    var packed_filter_ptr_dynamic = UnsafePointer[Scalar[type]].alloc(
        R * S * C * rounded_F_dynamic
    )
    var packed_filter_dynamic = LayoutTensor[type, layout_5d](
        packed_filter_ptr_dynamic,
        RuntimeLayout[layout_5d].row_major(
            Index(
                ceildiv(F, micro_kernel_f_size_default),
                R,
                S,
                C,
                micro_kernel_f_size_default,
            )
        ),
    )

    pack_filter(filter, packed_filter_dynamic, num_groups)

    # Conv attributes.
    comptime conv_attr_dynamic = ConvInfoStatic[2]()

    ConvDirectNHWC[
        input.layout,  # input shape
        layout_5d,  # filter shape
        layout_4d,  # output shape
        _,
        _,
        _,
        type,  # input type
        type,  # filter type
        type,  # output type
        True,
        conv_attr_dynamic,
    ].run(
        output_dynamic,
        input,
        packed_filter_dynamic,
        conv_shape,
    )

    comptime conv_attr_static = ConvInfoStatic[2](
        IntTuple(pad_h[0], pad_w[0], pad_h[1], pad_w[1]),
        IntTuple(stride[0], stride[1]),
        IntTuple(dilation[0], dilation[1]),
        num_groups,
    )

    comptime micro_kernel_shape = get_micro_kernel_shape[
        2, WO, F, conv_attr_static, simd_size
    ]()
    comptime micro_kernel_f_size = micro_kernel_shape[1] * simd_size
    comptime num_f_micro_tiles = ceildiv(F, micro_kernel_f_size)
    comptime rounded_F_static = num_f_micro_tiles * micro_kernel_f_size
    comptime packed_filter_layout = Layout.row_major(
        num_f_micro_tiles, R, S, C, micro_kernel_f_size
    )
    var packed_filter_ptr_static = UnsafePointer[Scalar[type]].alloc(
        R * S * C * rounded_F_static
    )
    var packed_filter_static = LayoutTensor[type, packed_filter_layout](
        packed_filter_ptr_static
    )

    pack_filter[simd_size, micro_kernel_f_size](
        filter,
        packed_filter_static,
        num_groups,
    )

    ConvDirectNHWC[
        Layout.row_major(N, H, W, C),
        packed_filter_layout,
        Layout.row_major(N, HO, WO, F),
        _,
        _,
        _,
        type,  # input type
        type,  # filter type
        type,  # output type
        True,
        conv_attr_static,
    ].run(
        output_static,
        input,
        packed_filter_static,
        conv_shape,
    )

    input_ptr.free()
    filter_ptr.free()
    packed_filter_ptr_dynamic.free()
    packed_filter_ptr_static.free()

    # Check results, return on the first failed comparison.
    for n, ho, wo, f in product(range(N), range(HO), range(WO), range(F)):
        if not isclose(
            output_dynamic[n, ho, wo, f],
            output_static[n, ho, wo, f],
            atol=1e-4,  # absolute error tolerance
            rtol=1e-5,  # relative error tolerance
        ):
            var expected = output_dynamic[n, ho, wo, f]
            var actual = output_static[n, ho, wo, f]
            print("Input shape NHWC: ", Index(N, H, W, C))
            print("filter shape RSCF: ", Index(R, S, C, F))
            print(
                "Failed at",
                Index(n, ho, wo, f),
                "expected",
                expected,
                "actual",
                actual,
                "rerr",
                abs(actual - expected) / abs(expected + 1e-10),
            )
            output_ptr_dynamic.free()
            output_ptr_static.free()
            return

    output_ptr_dynamic.free()
    output_ptr_static.free()

    # CHECK: Succeed
    print("Succeed")


def main():
    test[
        1,  # N
        14,  # H
        14,  # W
        256,  # C
        3,  # R
        3,  # S
        256,  # F
        Index(1, 1),  # stride
        Index(1, 1),  # dilation
        Index(1, 1),  # pad_h
        Index(1, 1),  # pad_w
    ]()
    test[
        1,  # N
        2,  # H
        2,  # W
        64,  # C
        3,  # R
        3,  # S
        64,  # F
        Index(2, 2),  # stride
        Index(1, 1),  # dilation
        Index(1, 1),  # pad_h
        Index(1, 1),  # pad_w
    ]()

    # Each test will build a specialization of the conv kernel.
    # Disable the following tests for now to monitor build time.

    # test[
    #     1,  # N
    #     56,  # H
    #     56,  # W
    #     64,  # C
    #     3,  # R
    #     3,  # S
    #     64,  # F
    #     Index(1, 1),  # stride
    #     Index(1, 1),  # dilation
    #     Index(1, 1),  # pad_h
    #     Index(1, 1),  # pad_w
    # ]()

    # test[
    #     1,  # N
    #     28,  # H
    #     28,  # W
    #     128,  # C
    #     3,  # R
    #     3,  # S
    #     128,  # F
    #     Index(1, 1),  # stride
    #     Index(1, 1),  # dilation
    #     Index(1, 1),  # pad_h
    #     Index(1, 1),  # pad_w
    # ]()

    # test[
    #     1,  # N
    #     7,  # H
    #     7,  # W
    #     512,  # C
    #     3,  # R
    #     3,  # S
    #     512,  # F
    #     Index(1, 1),  # stride
    #     Index(1, 1),  # dilation
    #     Index(1, 1),  # pad_h
    #     Index(1, 1),  # pad_w
    # ]()

    # test[
    #     1,  # N
    #     224,  # H
    #     224,  # W
    #     3,  # C
    #     7,  # R
    #     7,  # S
    #     64,  # F
    #     Index(2, 2),  # stride
    #     Index(1, 1),  # dilation
    #     Index(3, 3),  # pad_h
    #     Index(3, 3),  # pad_w
    # ]()

    # test[
    #     1,  # N
    #     56,  # H
    #     56,  # W
    #     128,  # C
    #     3,  # R
    #     3,  # S
    #     128,  # F
    #     Index(2, 2),  # stride
    #     Index(1, 1),  # dilation
    #     Index(1, 1),  # pad_h
    #     Index(1, 1),  # pad_w
    # ]()

    # test[
    #     1,  # N
    #     28,  # H
    #     28,  # W
    #     256,  # C
    #     3,  # R
    #     3,  # S
    #     256,  # F
    #     Index(2, 2),  # stride
    #     Index(1, 1),  # dilation
    #     Index(1, 1),  # pad_h
    #     Index(1, 1),  # pad_w
    # ]()

    # test[
    #     1,  # N
    #     14,  # H
    #     14,  # W
    #     512,  # C
    #     3,  # R
    #     3,  # S
    #     512,  # F
    #     Index(2, 2),  # stride
    #     Index(1, 1),  # dilation
    #     Index(1, 1),  # pad_h
    #     Index(1, 1),  # pad_w
    # ]()

    # test[
    #     19,  # N
    #     7,  # H
    #     7,  # W
    #     1,  # C
    #     3,  # R
    #     3,  # S
    #     16,  # F
    #     Index(1, 1),  # stride
    #     Index(1, 1),  # dilation
    #     Index(1, 1),  # pad_h
    #     Index(1, 1),  # pad_w
    # ]()

    # test[
    #     13,  # N
    #     14,  # H
    #     14,  # W
    #     2,  # C
    #     3,  # R
    #     3,  # S
    #     32,  # F
    #     Index(2, 2),  # stride
    #     Index(1, 1),  # dilation
    #     Index(1, 1),  # pad_h
    #     Index(1, 1),  # pad_w
    # ]()
